{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gengxuelong\n"
     ]
    }
   ],
   "source": [
    "print('gengxuelong')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "import yaml\n",
    "def load_dict_from_yaml(file_path: str):\n",
    "    with open(file_path, 'rt', encoding='utf-8') as f:\n",
    "        dict_1 = yaml.load(f, Loader=yaml.FullLoader)\n",
    "    return dict_1\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "import types\n",
    "args = load_dict_from_yaml('../args_recognize.yaml')\n",
    "args = types.SimpleNamespace(args)\n",
    "configs = load_dict_from_yaml('../configs.yaml')"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'attn_weight': 0.0, 'batch_size': 1, 'beam_size': 10, 'checkpoint': '/home/work_nfs7/xlgeng/bsmu_template/exp/salmonn_v8_lr5e_5/0.pt', 'config': '/home/work_nfs7/xlgeng/bsmu_template/exp/salmonn_v8_lr5e_5/train.yaml', 'context_bias_mode': '', 'context_graph_score': 0.0, 'context_list_path': '', 'ctc_weight': 0.5, 'data_type': 'raw', 'decoder_scale': 0.0, 'decoding_chunk_size': -1, 'gpu': 0, 'hlg': '', 'lm_scale': 0.0, 'modes': ['salmonn_decode'], 'num_decoding_left_chunks': -1, 'override_config': [], 'penalty': 0.0, 'r_decoder_scale': 0.0, 'result_dir': '/home/work_nfs7/xlgeng/bsmu_template/exp/salmonn_v8_lr5e_5/test_0pt/aishell1', 'reverse_weight': 0.0, 'search_ctc_weight': 1.0, 'search_transducer_weight': 0.0, 'simulate_streaming': False, 'test_data': '/home/work_nfs7/xlgeng/new_workspace/wenet_gxl_salmonn/examples/aishell/salmonn/data_list/test/aishell1/data.list', 'transducer_weight': 0.0, 'word': ''}\n"
     ]
    }
   ],
   "source": [
    "print(args)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'accum_grad': 14, 'ckpt_path': '/home/work_nfs7/xlgeng/bsmu_template/exp/salmonn_v7_lr5e_5/0.pt', 'dataset_conf': {'feats_type': 'raw_wav', 'filter_conf': {'max_length': 1670, 'min_length': 50}, 'resample_conf': {'resample_rate': 16000}, 'shuffle': True, 'shuffle_conf': {'shuffle_size': 1500}, 'sort': True, 'sort_conf': {'sort_size': 500}, 'batch_conf': {'batch_type': 'dynamic', 'batch_size': 5, 'max_frames_in_batch': 400000}}, 'frontend': 'hubert', 'frontend_conf': {'download_dir': './hub', 'frontend_conf': {'upstream': 'hubert_local', 'upstream_ckpt': '/home/work_nfs7/bsmu/SALMONN/resource/chinese_hubert_large.pt', 'upstream_model_config': None}, 'multilayer_feature': True}, 'grad_clip': 5, 'llama_model_generate_do_sample': True, 'llama_model_generate_length_penalty': 1.0, 'llama_model_generate_max_length': 80, 'llama_model_generate_min_length': 1, 'llama_model_generate_num_beams': 4, 'llama_model_generate_repetition_penalty': 1.0, 'llama_model_generate_temperature': 1.0, 'llama_model_generate_top_p': 0.9, 'llm_path': '/home/local_data/Atom-7B', 'load_epoch_ckpt': False, 'load_eval_ckpt': False, 'load_step_ckpt': True, 'log_interval': 500, 'lora_alpha': 32, 'lora_dropout': 0.1, 'lora_rank': 8, 'max_epoch': 15, 'optim': 'adamw', 'optim_conf': {'betas': [0.9, 0.99], 'eps': 1e-06, 'lr': 5e-05, 'weight_decay': 0.01}, 'prompt': 'Describe the speech.', 'scheduler': 'warmuplr', 'scheduler_conf': {'warmup_steps': 2000}, 'second_per_frame': 0.333333, 'second_stride': 0.333333, 'speech_qformer_layer': 2, 'speech_qformer_token_num': 1, 'use_lora': False, 'vicuna_low_resource': False, 'model': 'salmonn', 'tokenizer': 'llm', 'tokenizer_conf': {'llm_path': '/home/local_data/Atom-7B'}}\n"
     ]
    }
   ],
   "source": [
    "print(configs)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from wenet.dataset.dataset import Dataset\n",
    "import copy\n",
    "import logging\n",
    "import os\n",
    "from wenet.utils.init_tokenizer import init_tokenizer\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "                        format='%(asctime)s %(levelname)s %(message)s')\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)\n",
    "\n",
    "test_conf = copy.deepcopy(configs['dataset_conf'])\n",
    "\n",
    "test_conf['filter_conf']['max_length'] = 102400\n",
    "test_conf['filter_conf']['min_length'] = 0\n",
    "test_conf['filter_conf']['token_max_length'] = 102400\n",
    "test_conf['filter_conf']['token_min_length'] = 0\n",
    "test_conf['filter_conf']['max_output_input_ratio'] = 102400\n",
    "test_conf['filter_conf']['min_output_input_ratio'] = 0\n",
    "test_conf['speed_perturb'] = False\n",
    "test_conf['spec_aug'] = False\n",
    "test_conf['spec_sub'] = False\n",
    "test_conf['spec_trim'] = False\n",
    "test_conf['shuffle'] = False\n",
    "test_conf['sort'] = False\n",
    "if 'fbank_conf' in test_conf:\n",
    "    test_conf['fbank_conf']['dither'] = 0.0\n",
    "elif 'mfcc_conf' in test_conf:\n",
    "    test_conf['mfcc_conf']['dither'] = 0.0\n",
    "test_conf['batch_conf']['batch_type'] = \"static\"\n",
    "test_conf['batch_conf']['batch_size'] = args.batch_size\n",
    "\n",
    "tokenizer = init_tokenizer(configs)\n",
    "test_dataset = Dataset(args.data_type,\n",
    "                       args.test_data,\n",
    "                       tokenizer,\n",
    "                       test_conf,\n",
    "                       partition=False)\n",
    "\n",
    "test_data_loader = DataLoader(test_dataset, batch_size=None, num_workers=0)\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "for id, batch in enumerate(test_data_loader):\n",
    "    for k, v in batch.items():\n",
    "        if k == 'keys':\n",
    "            # print(k, v)\n",
    "            pass\n",
    "        else:\n",
    "            print(k, v.shape)\n",
    "    print('------------------------------------')\n",
    "    if id >= 1000:\n",
    "        break"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from wenet.utils.init_model import init_model\n",
    "\n",
    "# Init asr model from configs\n",
    "args.jit = False\n",
    "model, configs = init_model(args, configs)\n",
    "\n",
    "use_cuda = args.gpu >= 0 and torch.cuda.is_available()\n",
    "device = torch.device('cuda' if use_cuda else 'cpu')\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "print(model)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if \"salmonn_decode\" in args.modes:\n",
    "    print('decode mode: salmonn_decode')\n",
    "    result_file = os.path.join(args.result_dir, 'text_hyp')\n",
    "    with torch.no_grad(), open(result_file, 'w') as fout:\n",
    "        for batch_idx, batch in enumerate(test_data_loader):\n",
    "            sorted_keys, padded_feats, padding_labels, feats_lengths, label_lengths = batch\n",
    "            padded_feats = padded_feats.to(device)\n",
    "            feats_lengths = feats_lengths.to(device)\n",
    "            padding_labels = padding_labels.to(device)\n",
    "            label_lengths = label_lengths.to(device)\n",
    "            prompt = 'Describe the speech.'\n",
    "            if args.mode == 'salmonn_decode':\n",
    "                try:\n",
    "                    hyp = model.generate(padded_feats, feats_lengths, prompt)\n",
    "                except RuntimeError as e:\n",
    "                    logging.info(f'如下音频出现NaN：{sorted_keys}，error: {e}')\n",
    "\n",
    "            for i, key in enumerate(sorted_keys):\n",
    "                logging.info('{} {}'.format(key, hyp[0]))\n",
    "                fout.write('{} {}\\n'.format(key, hyp[0]))"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'types' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[6], line 3\u001B[0m\n\u001B[0;32m      1\u001B[0m configs \u001B[38;5;241m=\u001B[39m load_dict_from_yaml(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m../configs.yaml\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[0;32m      2\u001B[0m args \u001B[38;5;241m=\u001B[39m load_dict_from_yaml(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m../args_train.yaml\u001B[39m\u001B[38;5;124m'\u001B[39m)\n\u001B[1;32m----> 3\u001B[0m args \u001B[38;5;241m=\u001B[39m types\u001B[38;5;241m.\u001B[39mSimpleNamespace(args)\n\u001B[0;32m      4\u001B[0m \u001B[38;5;66;03m# init tokenizer\u001B[39;00m\n\u001B[0;32m      5\u001B[0m tokenizer \u001B[38;5;241m=\u001B[39m init_tokenizer(configs)\n",
      "\u001B[1;31mNameError\u001B[0m: name 'types' is not defined"
     ]
    }
   ],
   "source": [
    "from wenet.utils.train_utils import init_dataset_and_dataloader\n",
    "import types\n",
    "\n",
    "configs = load_dict_from_yaml(\"../configs.yaml\")\n",
    "args = load_dict_from_yaml('../args_train.yaml')\n",
    "args = types.SimpleNamespace(**args)\n",
    "# init tokenizer\n",
    "tokenizer = init_tokenizer(configs)\n",
    "\n",
    "train_dataset, cv_dataset, train_data_loader, cv_data_loader = \\\n",
    "    init_dataset_and_dataloader(args, configs, tokenizer)\n",
    "for id, batch in enumerate(train_data_loader):\n",
    "    for k, v in batch.items():\n",
    "        if k == 'keys':\n",
    "            # print(k, v)\n",
    "            pass\n",
    "        else:\n",
    "            print(k, v.shape)\n",
    "    print('------------------------------------')\n",
    "    if id >= 1000:\n",
    "        break\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
